import os
import glob
import h5py
import numpy as np
import cv2
import tensorflow as tf
from vgg16 import vgg16



def MSE(image,gt):
  
      width=image.shape[0]
      sumup=0
      for i in range(width):
           sumup=sumup+(image[i]-gt[i])*(image[i]-gt[i])
      MSE=sumup/width
      return MSE

def BG(image):
   images=[]
   MSEs=[]
   image=cv2.resize(image,(400,400))
   images.append(image)
   width=image.shape[0]
   height=image.shape[1]

   left_up=image[0:width*0.1,0:height*0.1]
   images.append(left_up)
   left_down=image[0:width*0.1,height*0.9:height]
   images.append(left_down)
   right_up=image[width*0.9:width,0:height*0.1]
   images.append(right_up)
   right_down=image[width*0.9:width,height*0.9:height]
   images.append(right_down)



def read_data(path):
     images = np.array(h5py.File(path, 'r').get('images'))
     labels = np.array(h5py.File(path,'r').get('labels'))
     return images, labels
     
def compare(x,y):
     stat_x = os.path.basename(x)
     stat_y = os.path.basename(y)
     if stat_x < stat_y:
          return -1
     elif stat_x > stat_y:
          return 1
     else:
          return 0

def prepare_data(sess, dataset):
     images_dir = os.path.join(os.getcwd(), dataset, 'images')
     labels_dir = os.path.join(os.getcwd(), dataset, 'labels')
     images = glob.glob(os.path.join(images_dir, '*.png'))
     labels = glob.glob(os.path.join(labels_dir, '*.png'))
     images.sort(compare)
     labels.sort(compare)
     return images, labels
     
def imread(path, is_singleChannel, img_sz):
     if is_singleChannel:
          im = cv2.imread(path)
          im = cv2.resize(im, (img_sz, img_sz))
          im_single = im[:,:,1]
          return im_single
     else:
          im = cv2.imread(path)
          im = cv2.resize(im, (img_sz, img_sz))
          return im

def make_data(sess, images, labels):
     savepath = os.path.join(os.getcwd(), 'train.h5')
     with h5py.File(savepath, 'w') as hf:
          hf.create_dataset('images', data=images)
          hf.create_dataset('labels', data=labels)
     
def input_setup(sess, img_sz):
     images, labels = prepare_data(sess, dataset='Train')
     input_seq = []
     label_seq = []

     for i in xrange(len(images)):
          image_ = imread(images[i], False, img_sz=img_sz)
          label_ = imread(labels[i], True, img_sz=img_sz)
          input_seq.append(image_)
          label_seq.append(label_)

     arrimages = np.asarray(input_seq)
     arrlabels = np.asarray(label_seq)
     make_data(sess, arrimages, arrlabels)


if __name__=='__main__':
     with tf.Session() as sess:
          input_setup(sess,400)
          data_dir = os.path.join(os.getcwd(), 'train.h5')
          train_images, train_labels = read_data(data_dir)
          print train_images.shape
          print train_labels.shape
          imgs = tf.placeholder(tf.float32, [None,224,224,3])
          vgg = vgg16(imgs, 'vgg16_weights.npz',sess)
          locs = []
          for i in range(train_images.shape[0]):
               probs=[]
               print i
               image = train_images[i,:,:,:]
               images=[]
               MSEs=[]
               images.append(image)
               width=image.shape[0]
               height=image.shape[1]
               left_up=image[0:width*0.1,0:height*0.1]
               images.append(left_up)
               left_down=image[0:width*0.1,height*0.9:height]
               images.append(left_down)
               right_up=image[width*0.9:width,0:height*0.1]
               images.append(right_up)
               right_down=image[width*0.9:width,height*0.9:height]
               images.append(right_down)
               for j in range (5):    
                    resized_img = cv2.resize(images[j],(224,224))
                    prob = sess.run(vgg.probs, feed_dict={vgg.imgs:[resized_img]})
                    probs.append(prob)
               image_prob=probs[0]
               MSEs=[MSE(probs[1],image_prob),MSE(probs[2],image_prob),MSE(probs[3],image_prob),MSE(probs[4],image_prob)]
               index=MSEs.index(min(MSEs))
               if index==0:
                   loc=[-1,-1]
               elif index==1:
                   loc=[-1,1]
               elif index==2:
                   loc=[1,-1]
               elif index==3:
                   loc=[1,1]
               print loc
               locs.append(loc)
          savepath = os.path.join(os.getcwd(), 'locs.h5')
          with h5py.File(savepath, 'w') as hf:
                hf.create_dataset('locs', data=locs)








     

